# Load AE model and Cluster Model for evaluation
# Author: YuFei Wang
# Date: 2024.11.18

# import libs
import os
import torch
from tqdm import tqdm 
import argparse
from argparse import ArgumentParser
import numpy as np

# import custom libs
from src.clustering_models.clusternet_modules.clusternetasmodel import ClusterNetModel
from AutoEncoder import DFConvAE
from src.utils import check_args, cluster_acc
from src.datasets import CustomDataset, TensorDatasetWrapper, transform_embeddings
from JointLearning import parse_minimal_args

class Evaluation:
    def __init__(self, args, input_dim=5000, latent_dim=20):
        self.args = args
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.init_k = 0
        
        self.train_loader, self.val_loader = self.Load_data()
        
        self.AE_model = DFConvAE(single_data_length=input_dim, latent_dim=latent_dim)
        self.Load_AE_model()
        
        state, init_k = self.Get_init_k()
        self.cluster_model = ClusterNetModel(args, input_dim=latent_dim, init_k=init_k)
        self.Load_Cluster_model(state, init_k)
    
    # 加载自编码器AutoEncoder模型
    def Load_AE_model(self,):
        state = torch.load("./saved_models/AE_models/AE_model.pth")
        try:
            self.AE_model.load_state_dict(state)
            print("Load AE model successfully!")
        except Exception as e:
            print("Load AE model failed! Error message: ", e)
            
    # 加载DeepDPM的聚类模型
    def Get_init_k(self,):
        state = torch.load("./saved_models/Cluster_models/Cluster_model.pth")
        init_k = state['cluster_net.class_fc2.weight'].shape[0] 
        return state, init_k
    
    def Load_Cluster_model(self, state, init_k):
        try:
            self.cluster_model.load_state_dict(state)
            print("Load Cluster model successfully!")
        except Exception as e:
            print("Load Cluster model failed! Error message: ", e)
    
    # 从CustomDataset类中加载数据
    def Load_data(self,):
        dataset_obj = CustomDataset(self.args, )
        train_loader, val_loader = dataset_obj.get_loaders()
        return train_loader, val_loader

    # 利用AE将数据转换为特征数据
    def Get_feature_step(self, datatype="train"):
        latent_x = []
        labels = []
        
        if datatype == "train":
            data_loader = self.train_loader
        elif datatype == "val":
            data_loader = self.val_loader
        else:
            raise NotImplementedError()
        
        for batch in tqdm(data_loader, desc="Getting data feature"):
            x, y = batch
            x = x.view(x.size(0), 1, -1)
            
            with torch.no_grad():
                latent_batch = self.AE_model.Encoder(x)
            
            latent_x.append(latent_batch)
            labels.append(y)
            
        latent_x = torch.cat(latent_x, dim=0)
        labels = torch.cat(labels, dim=0)
        print("Latent_x shape: ", latent_x.shape)
        print("Labels shape: ", labels.shape)
        
        # 转化过的特征需要经过transform_embeddings函数进行归一化
        latent_x = transform_embeddings("standard", latent_x)
        
        feature_set = TensorDatasetWrapper(latent_x, labels)
        feature_loader = torch.utils.data.DataLoader(feature_set, batch_size=self.args.batch_size, shuffle=False)
        
        return feature_loader
    
    # 得到数据的聚类分布
    def Get_assigment_step(self, feature_set):
        pred = []
        labels = []
        for batch in tqdm(feature_set, desc="Getting data assignment"):
            x, y = batch
            with torch.no_grad():
                assignment = self.cluster_model(x)
            pred.append(assignment)
            labels.append(y)
            
        pred = torch.cat(pred, dim=0)
        torch.save(pred, "pred.pt")
        labels = torch.cat(labels, dim=0)
        
        return pred, labels
    
    # 计算聚类准确率
    def Get_cluster_acc(self, pred, labels):
        pred = pred.argmax(dim=1).cpu().numpy()
        acc = np.round(cluster_acc(labels.cpu().numpy(), pred), 5)
        return acc
    
    # 评估模型
    def Evaluate(self,):
        train_feature_set = self.Get_feature_step(datatype="train")
        val_feature_set = self.Get_feature_step(datatype="val")
        
        train_pred, train_labels = self.Get_assigment_step(train_feature_set)
        val_pred, val_labels = self.Get_assigment_step(val_feature_set)
        
        train_acc = self.Get_cluster_acc(train_pred, train_labels)
        val_acc = self.Get_cluster_acc(val_pred, val_labels)
        
        print("Train Acc: {:.4f}, Val Acc: {:.4f}".format(train_acc, val_acc))
        
    # Function test
    def Test(self,):
        feature_set = self.Get_feature_step("val")
        
        pred, labels = self.Get_assigment_step(feature_set)
        print(pred.shape, labels.shape)

        acc = self.Get_cluster_acc(pred, labels)
        print("Acc: {:.4f}".format(acc))

# 测试函数，用于无法加载数据的时候的预备加载方案
def LoadPredFromFile():
    pred = torch.load("pred.pt")
    labels = torch.load("../SingleTabCluster/AutoEncoder/datasets/test_labels.pt")
    return pred, labels

      
if __name__ == "__main__":
    args = parse_minimal_args()
    EV = Evaluation(args)
    print("=" * 20)
    EV.Test()
    '''
    pred = torch.load("pred.pt")
    labels = torch.load("../SingleTabCluster/AutoEncoder/datasets/test_labels.pt")
    from collections import Counter
    print(Counter(pred.argmax(dim=1).cpu().numpy()))
    print(Counter(labels.cpu().numpy()))
    '''